import numpy as np
import open3d as o3d
import os
import copy
from pathlib import Path

############ 参数配置 ###########
# 文件路径
input_pcd = "../maps/outdoor_full.pcd"
output_pcd = "../maps/outdoor_final.pcd"
visualize_result = False  # 是否启用可视化验证

# 手动采样点（x, y）
sample_points = [
    [0, 0],
    [15, 0],
    [15, -20],
    [15, -40],
    [15, -70],
    [0, -70],
    [-20, -70],
    [-40, -60],
    [-30, 0],
    [-20, -10],
]
sample_radius = 1.0  # 采样半径（米）

def load_point_cloud(filepath):
    """加载点云文件并进行有效性检查，返回几何点云和张量点云"""
    if not Path(filepath).exists():
        raise FileNotFoundError(f"输入文件 {filepath} 未找到")

    # 使用t.io读取点云，保留所有属性
    t_pcd = o3d.t.io.read_point_cloud(filepath)
    if t_pcd.is_empty():
        raise ValueError("加载的点云文件为空")
    
    # 创建几何点云用于处理
    pcd_geom = o3d.geometry.PointCloud()
    pcd_geom.points = o3d.utility.Vector3dVector(t_pcd.point.positions.numpy())
    
    print(f"[1/6] 成功加载点云，包含 {len(pcd_geom.points)} 个点")
    return pcd_geom, t_pcd

def sample_ground_points(pcd, sample_points, radius):
    """使用KDTree进行高效邻域采样"""
    print("[2/6] 开始点采样...")
    kdtree = o3d.geometry.KDTreeFlann(pcd)
    points_array = np.asarray(pcd.points)
    sampled_points = []

    for idx, (x, y) in enumerate(sample_points):
        # 搜索球形邻域内的点
        [k, idxs, _] = kdtree.search_hybrid_vector_3d(
            query=[x, y, 0],  # Z轴方向忽略
            radius=radius,
            max_nn=100,  # 限制最大邻近点数
        )

        if k < 3:
            print(f"警告: 采样点 ({x}, {y}) 邻域内点数不足 (仅 {k} 个)，已跳过")
            continue

        # 计算邻域点的平均高度
        z_avg = np.mean(points_array[idxs, 2])
        sampled_points.append([x, y, z_avg])
        print(
            f"采样点 {idx+1}/{len(sample_points)} 处理完成，找到 {k} 个邻近点", end="\r"
        )

    print(f"\n共获取 {len(sampled_points)} 个有效采样点")
    return np.array(sampled_points)

def fit_plane_svd(points):
    """使用SVD进行平面拟合"""
    print("[3/6] 正在拟合平面...")
    if len(points) < 3:
        raise ValueError("采样点不足3个，无法拟合平面")

    centroid = np.mean(points, axis=0)
    _, _, Vt = np.linalg.svd(points - centroid)
    normal_vector = Vt[2, :]

    # 确保法向量朝向Z轴正方向
    if normal_vector[2] < 0:
        normal_vector = -normal_vector

    print(f"拟合平面法向量: {normal_vector}")
    return normal_vector

def compute_rotation_matrix(source):
    """计算对齐到Z轴的旋转矩阵"""
    print("[4/6] 计算旋转矩阵...")
    target = np.array([0, 0, 1], dtype=np.float64)
    source = source / np.linalg.norm(source)

    # 使用Open3D内置方法简化计算
    if np.allclose(source, target):
        return np.eye(3)

    axis = np.cross(source, target)
    axis = axis / np.linalg.norm(axis)
    angle = np.arccos(np.dot(source, target))

    # 罗德里格斯公式优化实现
    K = np.array(
        [[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]]
    )
    return np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)

def align_point_cloud(pcd, rotation_matrix):
    """应用旋转矩阵进行点云调平"""
    print("[5/6] 正在应用旋转...")
    # 深拷贝点云以保留原始数据
    aligned_pcd = copy.deepcopy(pcd)
    aligned_pcd.rotate(rotation_matrix, center=(0, 0, 0))
    return aligned_pcd

def save_point_cloud(t_pcd, filepath):
    """保存点云并创建输出目录"""
    print("[6/6] 保存结果...")
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    o3d.t.io.write_point_cloud(filepath, t_pcd)
    print(f"调平完成，文件保存至 {filepath}")

def visualize_results(original, transformed):
    """可视化原始与调平后的点云对比"""
    if not visualize_result:
        return

    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="点云调平对比")

    # 创建几何体副本并设置不同颜色
    original.paint_uniform_color([1, 0, 0])  # 红色表示原始点云
    transformed.paint_uniform_color([0, 1, 0])  # 绿色表示调平后的点云

    vis.add_geometry(original)
    vis.add_geometry(transformed)

    render_option = vis.get_render_option()
    render_option.point_size = 1.0
    render_option.background_color = np.asarray([0, 0, 0])

    vis.run()
    vis.destroy_window()

def main():
    try:
        # 0. 输出参数
        print("点云调程序启动...")
        print(f"输入点云: {input_pcd}")
        print(f"输出点云: {output_pcd}")
        
        # 1. 加载点云（同时获取几何点云和张量点云）
        pcd_geom, t_pcd = load_point_cloud(input_pcd)

        # 2. 采样地面点
        sampled_points = sample_ground_points(pcd_geom, sample_points, sample_radius)

        # 3. 平面拟合
        normal_vector = fit_plane_svd(sampled_points)

        # 4. 计算旋转矩阵
        rotation_matrix = compute_rotation_matrix(normal_vector)

        # 5. 应用旋转
        aligned_geom = align_point_cloud(pcd_geom, rotation_matrix)
        
        # 5.1 更新张量点云的点坐标（保留所有原始属性）
        rotated_points = np.asarray(aligned_geom.points)
        t_pcd.point.positions = o3d.core.Tensor(rotated_points, dtype=o3d.core.Dtype.Float32)

        # 6. 保存结果（保留所有字段）
        save_point_cloud(t_pcd, output_pcd)

        # 7. 可视化（可选）
        if visualize_result:
            visualize_results(pcd_geom, aligned_geom)

    except Exception as e:
        print(f"错误: {str(e)}")
        raise

if __name__ == "__main__":
    main()